import socket
import time

# ================= config =================
HOST = '0.0.0.0'
PORT = 50005     
CHUNK_SIZE = 1024 
# =======================================

def main():
    # 1. create UDP Socket
    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    sock.bind((HOST, PORT))
    
    print(f"🔄 Echo test server started")
    print(f"👉 Listening the port {PORT}，waiting Pico to send audio...")
    
    # Buffer to store incoming audio data
    audio_buffer = bytearray()
    client_addr = None

    while True:
        try:
            # receive data from Pico
            data, addr = sock.recvfrom(2048)
            
            #if no client address yet, set it
            if client_addr is None:
                client_addr = addr

            # ---logic---
            if data == b'START':
                print(f"\n[1] Begin the recieveing (from {addr})...")
                audio_buffer = bytearray() # clear buffer
                
            elif data == b'STOP':
                print(f"[2] Recording finished with {len(audio_buffer)} bytes。")
                print(f"[3] Sending back the audio (Echo)...")
                
                # --- Echo ---
                total_sent = 0
                
                # cut the buffer and send in chunks
                for i in range(0, len(audio_buffer), CHUNK_SIZE):
                    chunk = audio_buffer[i : i + CHUNK_SIZE]
                    sock.sendto(chunk, addr)
                    total_sent += len(chunk)
                    
                    # tiny delay to avoid overwhelming the network
                    time.sleep(0.035) 
                
                # the end signal
                sock.sendto(b'END_OF_STREAM', addr)
                print(f"✅ Finsihed, waiting for next one...")
                
            else:
                # save incoming audio data
                audio_buffer.extend(data)
                # print a dot every 10KB received
                if len(audio_buffer) % 10240 == 0:
                    print(".", end="", flush=True)

        except KeyboardInterrupt:
            print("\nServer stopped by user")
            break
        except Exception as e:
            print(f"Wrong: {e}")

if __name__ == "__main__":
    main()